import pandas as pd
import os
from typing import List, Dict, Optional
from datetime import datetime, timedelta

def run_backtest_all_market(
    start_date: str,
    end_date: str,
    stop_loss_pct: float,
    take_profit_pct: float,
    initial_cash: float,
    max_holdings: int
) -> Dict:
    """
    Performs a full-market stock backtest, selecting top stocks daily based on Sharpe ratio.

    Args:
        start_date (str): Start date for backtesting (YYYY-MM-DD).
        end_date (str): End date for backtesting (YYYY-MM-DD).
        stop_loss_pct (float): Stop-loss percentage (e.g., 0.1 for 10% loss).
        take_profit_pct (float): Take-profit percentage (e.g., 0.2 for 20% gain).
        initial_cash (float): Initial cash for the portfolio.
        max_holdings (int): Maximum number of stocks to hold at any time.

    Returns:
        dict: Contains transaction records, portfolio history, holdings, and summary.
    """
    # --- Configuration ---
    data_file = 'job/data/trade_indicators.parquet'
    output_dir = 'data/backtest_results'
    output_transactions_file = os.path.join(output_dir, f'backtest_all_market_transactions.csv')
    output_portfolio_history_file = os.path.join(output_dir, f'backtest_all_market_portfolio_history.csv')
    sharpe_start_date = '2023-01-01'

    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    print(f"Output directory '{output_dir}' ensured to exist.")

    # --- Data Loading and Preprocessing ---
    try:
        df = pd.read_parquet(data_file)
        print(f"Successfully loaded data from '{data_file}'.")
    except FileNotFoundError:
        print(f"Error: Data file '{data_file}' not found.")
        return {
            "message": f"Data file '{data_file}' not found.",
            "transactions": [],
            "portfolio_history": [],
            "current_holdings": [],
            "summary": {"error": f"Data file '{data_file}' not found."}
        }
    except Exception as e:
        print(f"Error loading Parquet file: {e}")
        return {
            "message": "Error loading data.",
            "transactions": [],
            "portfolio_history": [],
            "current_holdings": [],
            "summary": {"error": f"Error loading Parquet file: {e}"}
        }

    # Handle date column
    if df.index.name == 'date':
        df = df.reset_index()
    if 'datekey' in df.columns:
        df['date'] = pd.to_datetime(df['datekey'].astype(str).str.replace('d', ''), format='%Y%m%d')
    else:
        print("Error: 'datekey' or 'date' column not found.")
        return {
            "message": "Missing date column.",
            "transactions": [],
            "portfolio_history": [],
            "current_holdings": [],
            "summary": {"error": "Missing 'datekey' or 'date' column in data."}
        }

    df['code'] = df['code'].astype(str)
    df = df.sort_values(by=['date', 'code']).reset_index(drop=True)

    # Filter data by date range
    filtered_df = df[(df['date'] >= pd.to_datetime(start_date)) &
                    (df['date'] <= pd.to_datetime(end_date))].copy()

    if filtered_df.empty:
        print(f"No data found in date range {start_date} to {end_date}.")
        return {
            "message": f"No data found in the specified date range.",
            "transactions": [],
            "portfolio_history": [],
            "current_holdings": [],
            "summary": {
                "start_date": start_date,
                "end_date": end_date,
                "initial_cash": round(initial_cash, 2),
                "final_cash": round(initial_cash, 2),
                "final_portfolio_value": round(initial_cash, 2),
                "sharpe_ratio_sum": None
            }
        }

    # Derive 'action' column if not present
    if 'action' not in filtered_df.columns:
        raise ValueError("The 'action' column is required in the data.")

    filtered_df['next_day_action'] = filtered_df.groupby('code')['action'].shift(1)
    filtered_df['ma_5_close'] = filtered_df.groupby('code')['close'].rolling(window=5, min_periods=1).mean().shift(1).reset_index(level=0, drop=True)

    # --- Backtesting Simulation ---
    cash = float(initial_cash)
    holdings = {}  # Dict to track {stock_code: {'shares': int, 'buy_price': float, 'buy_date': str}}
    transactions = []
    portfolio_history = []
    total_sharpe_sum = 0.0

    unique_dates = filtered_df['date'].unique()
    for current_date in unique_dates:
        current_date_str = pd.Timestamp(current_date).strftime('%Y-%m-%d')
        print(f"\nProcessing date: {current_date_str}")

        # Get previous day's data for buy/sell signals
        prev_date = pd.Timestamp(current_date) - timedelta(days=1)
        prev_day_df = df[df['date'] == prev_date].copy()

        # Calculate cumulative Sharpe ratio for each stock up to prev_date
        sharpe_df = df[(df['date'] >= pd.to_datetime(sharpe_start_date)) &
                      (df['date'] <= prev_date)].copy()
        if 'sharpe_ratio' in sharpe_df.columns:
            sharpe_sum = sharpe_df.groupby('code')['sharpe_ratio'].sum().reset_index()
        else:
            print(f"Warning: 'sharpe_ratio' column not found for {current_date_str}.")
            sharpe_sum = pd.DataFrame({'code': [], 'sharpe_ratio': []})

        # Get stocks with 'buy' signal for the current day
        buy_candidates = prev_day_df[prev_day_df['action'] == 'buy'][['code']].copy()
        if buy_candidates.empty:
            print(f"No buy signals for {current_date_str}.")
        else:
            buy_candidates = buy_candidates.merge(sharpe_sum, on='code', how='left')
            buy_candidates['sharpe_ratio'] = buy_candidates['sharpe_ratio'].fillna(0)
            buy_candidates = buy_candidates.sort_values(by='sharpe_ratio', ascending=False)

        # Get current day's data
        current_day_df = filtered_df[filtered_df['date'] == current_date].copy()

        # --- Sell Logic (Stop-loss, Take-profit, or Sell Signal) ---
        for stock_code in list(holdings.keys()):
            stock_data = current_day_df[current_day_df['code'] == stock_code]
            if stock_data.empty:
                continue

            stock_data = stock_data.iloc[0]
            current_close = stock_data['close']
            current_open = stock_data['open']
            current_high = stock_data['high']
            current_low = stock_data['low']
            action = stock_data['next_day_action']
            holding = holdings[stock_code]
            buy_price = holding['buy_price']
            shares_held = holding['shares']

            if pd.isna(current_close) or pd.isna(current_open) or pd.isna(current_high) or pd.isna(current_low):
                print(f"Skipping {stock_code} on {current_date_str}: Missing price data.")
                continue

            # Check stop-loss or take-profit
            price_change = (current_open - buy_price) / buy_price
            should_sell = (
                price_change <= -stop_loss_pct or
                price_change >= take_profit_pct or
                action == 'sell'
            )

            if should_sell and shares_held > 0:
                sell_price = current_open
                sell_amount = shares_held * sell_price
                cash += sell_amount
                print(f"{current_date_str}: SOLD {shares_held} shares of {stock_code} at {sell_price:.2f}. "
                      f"Proceeds: {sell_amount:.2f}. Reason: {'Stop-loss' if price_change <= -stop_loss_pct else 'Take-profit' if price_change >= take_profit_pct else 'Sell signal'}.")
                transactions.append({
                    'date': current_date_str,
                    'type': 'SELL',
                    'stock_code': stock_code,
                    'price': round(sell_price, 2),
                    'shares': shares_held,
                    'amount': round(sell_amount, 2),
                    'cash_after_transaction': round(cash, 2),
                    'shares_after_transaction': 0
                })
                del holdings[stock_code]

        # Record portfolio value before buys
        portfolio_value = cash
        for stock_code, holding in holdings.items():
            stock_data = current_day_df[current_day_df['code'] == stock_code]
            if not stock_data.empty:
                portfolio_value += holding['shares'] * stock_data.iloc[0]['close']

        portfolio_history.append({
            'date': current_date_str,
            'portfolio_value': round(portfolio_value, 2),
            'cash': round(cash, 2),
            'holdings': {k: {'shares': v['shares'], 'buy_price': v['buy_price'], 'buy_date': v['buy_date']} for k, v in holdings.items()}
        })

        # --- Buy Logic ---
        available_slots = max_holdings - len(holdings)
        if available_slots > 0 and not buy_candidates.empty:
            cash_per_stock = cash / available_slots if available_slots > 0 else 0
            for _, candidate in buy_candidates.head(available_slots).iterrows():
                stock_code = candidate['code']
                if stock_code in holdings:
                    continue

                stock_data = current_day_df[current_day_df['code'] == stock_code]
                if stock_data.empty:
                    continue

                stock_data = stock_data.iloc[0]
                current_open = stock_data['open']
                current_high = stock_data['high']
                current_low = stock_data['low']
                buy_price_ma = stock_data['ma_5_close']

                if pd.isna(current_open) or pd.isna(current_high) or pd.isna(current_low) or pd.isna(buy_price_ma):
                    print(f"Skipping {stock_code} on {current_date_str}: Missing price or MA data.")
                    continue

                buy_price = buy_price_ma if current_low <= buy_price_ma <= current_high else current_open
                if cash_per_stock >= buy_price:
                    shares_to_buy = int(cash_per_stock / buy_price)
                    if shares_to_buy > 0:
                        buy_amount = shares_to_buy * buy_price
                        cash -= buy_amount
                        holdings[stock_code] = {
                            'shares': shares_to_buy,
                            'buy_price': buy_price,
                            'buy_date': current_date_str
                        }
                        print(f"{current_date_str}: BOUGHT {shares_to_buy} shares of {stock_code} at {buy_price:.2f}. "
                              f"Cost: {buy_amount:.2f}.")
                        transactions.append({
                            'date': current_date_str,
                            'type': 'BUY',
                            'stock_code': stock_code,
                            'price': round(buy_price, 2),
                            'shares': shares_to_buy,
                            'amount': round(buy_amount, 2),
                            'cash_after_transaction': round(cash, 2),
                            'shares_after_transaction': shares_to_buy
                        })

        # Update total Sharpe ratio
        day_sharpe = current_day_df['sharpe_ratio'].sum() if 'sharpe_ratio' in current_day_df.columns else 0
        total_sharpe_sum += day_sharpe

    # --- Final Summary ---
    final_portfolio_value = cash
    current_holdings = []
    for stock_code, holding in holdings.items():
        stock_data = filtered_df[(filtered_df['code'] == stock_code) & (filtered_df['date'] == unique_dates[-1])]
        current_price = stock_data.iloc[0]['close'] if not stock_data.empty else 0
        if not pd.isna(current_price):
            final_portfolio_value += holding['shares'] * current_price
        current_holdings.append({
            'stock_code': stock_code,
            'shares': holding['shares'],
            'buy_price': round(holding['buy_price'], 2),
            'buy_date': holding['buy_date'],
            'current_price': round(current_price, 2) if not pd.isna(current_price) else None,
            'unrealized_pnl': round((current_price - holding['buy_price']) * holding['shares'], 2) if not pd.isna(current_price) else None
        })

    print(f"\n--- Backtest Summary ---")
    print(f"Initial Cash: {initial_cash:.2f}")
    print(f"Final Cash: {cash:.2f}")
    print(f"Final Portfolio Value: {final_portfolio_value:.2f}")
    print(f"Total Sharpe Ratio Sum: {total_sharpe_sum:.4f}")
    print(f"Current Holdings: {len(holdings)} stocks")

    # Save to CSVs
    if transactions:
        pd.DataFrame(transactions).to_csv(output_transactions_file, index=False)
        print(f"Transactions saved to: {output_transactions_file}")
    if portfolio_history:
        pd.DataFrame(portfolio_history).to_csv(output_portfolio_history_file, index=False)
        print(f"Portfolio history saved to: {output_portfolio_history_file}")

    return {
        "message": "Backtest completed successfully.",
        "transactions": transactions,
        "portfolio_history": portfolio_history,
        "current_holdings": current_holdings,
        "summary": {
            "start_date": start_date,
            "end_date": end_date,
            "initial_cash": round(initial_cash, 2),
            "final_cash": round(cash, 2),
            "final_portfolio_value": round(final_portfolio_value, 2),
            "sharpe_ratio_sum": round(total_sharpe_sum, 4) if total_sharpe_sum != 0 else None
        }
    }

if __name__ == '__main__':
    result = run_backtest_all_market(
        start_date="2023-01-01",
        end_date="2025-06-10",
        stop_loss_pct=0.1,
        take_profit_pct=0.2,
        initial_cash=100000,
        max_holdings=5
    )
    print("\n--- Backtest Results ---")
    print(result)